import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def load_cifar10_data(batch_size=128):
    """
    加载CIFAR10数据集并进行数据增强
    
    参数:
        batch_size (int): 每个批次的样本数
        
    返回:
        train_loader (DataLoader): 训练集数据加载器
        test_loader (DataLoader): 测试集数据加载器
        class_names (tuple): 类别名称元组
    """
    # 定义训练集的数据增强和预处理
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),  # 随机裁剪
        transforms.RandomHorizontalFlip(),     # 随机水平翻转
        transforms.RandomRotation(15),         # 随机旋转
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色抖动
        transforms.ToTensor(),                 # 转换为张量
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))  # 标准化
    ])

    # 测试集的预处理(不进行数据增强)
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])

    # 加载训练集
    train_dataset = datasets.CIFAR10(
        root='./data', 
        train=True,
        download=True, 
        transform=train_transform
    )
    
    # 使用数据并行加载器(如果可用)
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True if torch.cuda.is_available() else False
    )

    # 加载测试集
    test_dataset = datasets.CIFAR10(
        root='./data', 
        train=False,
        download=True, 
        transform=test_transform
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True if torch.cuda.is_available() else False
    )

    # CIFAR10的类别名称
    class_names = ('plane', 'car', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck')

    return train_loader, test_loader, class_names